﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using Ast.Core.Pe;
using Ast.Core.Pe.Definitions;
using Ast.Framework;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface IPeService : IDependency
    {
        PeModule Load(string filePath);
        PeModule Load(Stream fileStream, bool closeStream);
    }

    public class PeService : IPeService
    {
        public PeService()
        {
            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public PeModule Load(string filePath)
        {
            if (!File.Exists(filePath))
            {
                Logger.ErrorFormat("file not found at {0}", filePath);
                return null;
            }
            using (var fileStream = File.Open(filePath, FileMode.Open, FileAccess.Read, FileShare.Read))
            {
                return Load(fileStream, true);
            }
        }

        public PeModule Load(Stream fileStream, bool closeStream)
        {
            var peModule = new PeModule();
            using (var binaryReader = new PeBinaryReader(fileStream, Encoding.UTF8, !closeStream))
            {
                DosHeader dosHeader;
                if (ReadDosHeader(binaryReader, out dosHeader))
                {
                    peModule.DosHeader = dosHeader;
                }
                else
                {
                    Logger.Error("Error occurred in ReadDosHeader");
                    return null;
                }

                NewExeFileHeader newExeFileHeader;
                if (ReadNewExeFileHeader(binaryReader, dosHeader, out newExeFileHeader))
                {
                    peModule.NewExeFileHeader = newExeFileHeader;
                }
                else
                {
                    Logger.Error("Error occurred in ReadNewExeFileHeader");
                    return null;
                }

                peModule.OptionalHeader = ReadOptionalHeader(binaryReader, dosHeader, newExeFileHeader);
                if (peModule.ModuleType == PeModuleType.Unknown)
                {
                    Logger.Error("Error occurred in ReadOptionalHeader");
                    return null;
                }

                peModule.SectionHeaders = ReadSectionHeaders(binaryReader, dosHeader, newExeFileHeader);
                if (peModule.SectionHeaders == null || peModule.SectionHeaders.Length == 0)
                {
                    Logger.Error("Error occurred in ReadSectionHeaders");
                    return null;
                }

                peModule.SectionContents = new Dictionary<SectionHeader, byte[]>(peModule.SectionHeaders.Length);
                ReadSectionContents(binaryReader, peModule);

                peModule.RelocatedBlocks = new List<RelocatedBlock>();
                ReadRelocatedBlocks(binaryReader, peModule);

                peModule.ImportedModules = new List<ImportedModule>();
                ReadImportedModules(binaryReader, peModule);
                ReadDeplayImportedModules(binaryReader, peModule);

                peModule.ExportedFunctions = new List<ExportedFunction>();
                ReadExportedFunctions(binaryReader, peModule);
            }
            return peModule;
        }

        private void ReadRelocatedBlocks(PeBinaryReader reader, PeModule peModule)
        {
            if (!peModule.RelocationDirectoryEntry.HasValue) return;

            var directoryEntry = peModule.RelocationDirectoryEntry.Value;
            var relocationDirectoryDescriptorPosition = peModule.RvaToFileOffset(directoryEntry.VirtualAddress);
            if (!relocationDirectoryDescriptorPosition.HasValue) return;

            reader.SetPosition(relocationDirectoryDescriptorPosition.Value);
            var relocationDirectoryDescriptor = reader.ReadAs<RelocationDirectoryDescriptor>();
            while (relocationDirectoryDescriptor.VirtualAddress != 0)
            {
                var relocationTypeOffsets = reader.ReadArrayAs<RelocationTypeOffset>(relocationDirectoryDescriptor.Count);
                var relocatedBlock = new RelocatedBlock(relocationDirectoryDescriptor.VirtualAddress, relocationDirectoryDescriptor.SizeOfBlock)
                {
                    TypeOffsets = relocationTypeOffsets
                };
                peModule.RelocatedBlocks.Add(relocatedBlock);

                relocationDirectoryDescriptor = reader.ReadAs<RelocationDirectoryDescriptor>();
            }
        }

        private void ReadSectionContents(PeBinaryReader reader, PeModule peModule)
        {
            foreach (var sectionHeader in peModule.SectionHeaders)
            {
                var size = Math.Max(sectionHeader.PhysicalAddressOrVirtualSize, sectionHeader.SizeOfRawData);
                reader.SetPosition(sectionHeader.PointerToRawData);
                var content = reader.ReadBytes((int) size);
                peModule.SectionContents.Add(sectionHeader, content);
            }
        }

        private bool ReadDosHeader(PeBinaryReader reader, out DosHeader dosHeader)
        {
            reader.SetPosition(0);
            dosHeader = reader.ReadAs<DosHeader>();
            return dosHeader.IsValid;
        }

        private bool ReadNewExeFileHeader(PeBinaryReader reader, DosHeader dosHeader, out NewExeFileHeader newExeFileHeader)
        {
            reader.SetPosition(dosHeader.NewExeHeaderOffset);
            newExeFileHeader = reader.ReadAs<NewExeFileHeader>();
            return newExeFileHeader.IsValid;
        }

        private IOptionalHeader ReadOptionalHeader(PeBinaryReader reader, DosHeader dosHeader, NewExeFileHeader newExeFileHeader)
        {
            var optionalHeaderPosition = dosHeader.NewExeHeaderOffset + Marshal.SizeOf<NewExeFileHeader>();
            reader.SetPosition(optionalHeaderPosition);
            if (newExeFileHeader.SizeOfOptionalHeader == OptionalHeader32.StructSize)
            {
                return reader.ReadAs<OptionalHeader32>();
            }
            if (newExeFileHeader.SizeOfOptionalHeader == OptionalHeader64.StructSize)
            {
                return reader.ReadAs<OptionalHeader64>();
            }
            return null;
        }

        private SectionHeader[] ReadSectionHeaders(PeBinaryReader reader, DosHeader dosHeader, NewExeFileHeader newExeFileHeader)
        {
            var sectionHeaderPosition = dosHeader.NewExeHeaderOffset + Marshal.SizeOf<NewExeFileHeader>() + newExeFileHeader.SizeOfOptionalHeader;
            reader.SetPosition(sectionHeaderPosition);
            return reader.ReadArrayAs<SectionHeader>(newExeFileHeader.NumberOfSections);
        }

        private void ReadImportedModules(PeBinaryReader reader, PeModule peModule)
        {
            if (!peModule.ImportDirectoryEntry.HasValue) return;

            var directoryEntry = peModule.ImportDirectoryEntry.Value;
            var importDirectoryDescriptorPosition = peModule.RvaToFileOffset(directoryEntry.VirtualAddress);
            if (!importDirectoryDescriptorPosition.HasValue) return;

            reader.SetPosition(importDirectoryDescriptorPosition.Value);
            var importDirectoryDescriptors = reader.ReadArrayAs<ImportDirectoryDescriptor>();
            foreach (var importDirectoryDescriptor in importDirectoryDescriptors)
            {
                var moduleNamePosition = peModule.RvaToFileOffset(importDirectoryDescriptor.Name);
                if (!moduleNamePosition.HasValue) continue;

                reader.SetPosition(moduleNamePosition.Value);
                var moduleName = reader.ReadAsString();
                var importedModule = new ImportedModule(moduleName, importDirectoryDescriptor.TimeDateStamp, false);
                ReadImportedFunctions(reader, peModule, importedModule, importDirectoryDescriptor.OriginalFirstThunk, importDirectoryDescriptor.FirstThunk);
                peModule.ImportedModules.Add(importedModule);
            }
        }

        private void ReadDeplayImportedModules(PeBinaryReader reader, PeModule peModule)
        {
            if (!peModule.DelayLoadImportDirectoryEntry.HasValue) return;

            var directoryEntry = peModule.DelayLoadImportDirectoryEntry.Value;
            var delayLoadDirectoryDescriptorPosition = peModule.RvaToFileOffset(directoryEntry.VirtualAddress);
            if (!delayLoadDirectoryDescriptorPosition.HasValue) return;

            reader.SetPosition(delayLoadDirectoryDescriptorPosition.Value);
            var delayLoadDirectoryDescriptors = reader.ReadArrayAs<DelayLoadDirectoryDescriptor>();
            foreach (var delayLoadDirectoryDescriptor in delayLoadDirectoryDescriptors)
            {
                var moduleNamePosition = peModule.RvaToFileOffset(delayLoadDirectoryDescriptor.Name);
                if (!moduleNamePosition.HasValue) continue;

                reader.SetPosition(moduleNamePosition.Value);
                var moduleName = reader.ReadAsString();
                var importedModule = new ImportedModule(moduleName, (int) delayLoadDirectoryDescriptor.TimeStamp, true);
                ReadImportedFunctions(reader, peModule, importedModule, delayLoadDirectoryDescriptor.DelayImportNameTable, delayLoadDirectoryDescriptor.DelayImportAddressTable);
                peModule.ImportedModules.Add(importedModule);
            }
        }

        private void ReadExportedFunctions(PeBinaryReader reader, PeModule peModule)
        {
            if (!peModule.ExportDirectoryEntry.HasValue) return;

            var directoryEntry = peModule.ExportDirectoryEntry.Value;
            var delayLoadDirectoryDescriptorPosition = peModule.RvaToFileOffset(directoryEntry.VirtualAddress);
            if (!delayLoadDirectoryDescriptorPosition.HasValue) return;

            reader.SetPosition(delayLoadDirectoryDescriptorPosition.Value);
            var exportDirectoryHeader = reader.ReadAs<ExportDirectoryHeader>();

            var functionRvasPosition = peModule.RvaToFileOffset(exportDirectoryHeader.AddressOfFunctions);
            if (!functionRvasPosition.HasValue) return;
            var nameOffsetsPosition = peModule.RvaToFileOffset(exportDirectoryHeader.AddressOfNames);
            if (!nameOffsetsPosition.HasValue) return;
            var nameOrdinalsPosition = peModule.RvaToFileOffset(exportDirectoryHeader.AddressOfNameOrdinals);
            if (!nameOrdinalsPosition.HasValue) return;

            reader.SetPosition(functionRvasPosition.Value);
            var functionRvas = reader.ReadArrayAs<uint>((int) exportDirectoryHeader.NumberOfFunctions);
            reader.SetPosition(nameOffsetsPosition.Value);
            var nameOffsets = reader.ReadArrayAs<uint>((int) exportDirectoryHeader.NumberOfNames);
            reader.SetPosition(nameOrdinalsPosition.Value);
            var nameOrdinals = reader.ReadArrayAs<ushort>((int) exportDirectoryHeader.NumberOfNames);

            var nameOrdinalTable = new Dictionary<int, string>();
            for (var i = 0; i < exportDirectoryHeader.NumberOfNames; i++)
            {
                var ordinal = nameOrdinals[i];
                var nameOffset = nameOffsets[i];
                var namePosition = peModule.RvaToFileOffset(nameOffset);
                if (!namePosition.HasValue) continue;

                reader.SetPosition(namePosition.Value);
                var name = reader.ReadAsString();
                nameOrdinalTable.Add(ordinal, name);
            }

            for (var i = 0; i < exportDirectoryHeader.NumberOfFunctions; i++)
            {
                var functionAddress = functionRvas[i];
                if (functionAddress == 0) continue;

                string name;
                nameOrdinalTable.TryGetValue(i, out name);

                string forwardModuleName = null;
                string forwardFunctionName = null;
                if (directoryEntry.Contains(functionAddress))
                {
                    var forwarderPosition = peModule.RvaToFileOffset(functionAddress);
                    if (!forwarderPosition.HasValue) continue;

                    reader.SetPosition(forwarderPosition.Value);
                    var forwarder = reader.ReadAsString();
                    var iExclaim = forwarder.IndexOf('.');
                    if (iExclaim > 0)
                    {
                        forwardModuleName = forwarder.Substring(0, iExclaim);
                        forwardFunctionName = forwarder.Substring(iExclaim + 1);
                    }
                    else
                    {
                        forwardFunctionName = forwarder;
                    }
                }

                var exportedFunction = new ExportedFunction(functionAddress, i + exportDirectoryHeader.Base, name, forwardModuleName, forwardFunctionName);
                peModule.ExportedFunctions.Add(exportedFunction);
            }
        }

        private void ReadImportedFunctions(PeBinaryReader reader, PeModule peModule, ImportedModule importedModule, uint intRva, uint iatRva)
        {
            var lookupTableRva = intRva != 0 ? intRva : iatRva;
            if (lookupTableRva == 0) return;

            var lookupTablePosition = peModule.RvaToFileOffset(lookupTableRva);
            if (!lookupTablePosition.HasValue) return;

            IImportFunctionDescriptor[] functionDescriptors = null;
            reader.SetPosition(lookupTablePosition.Value);
            if (peModule.ModuleType == PeModuleType.Pe32)
            {
                functionDescriptors = reader.ReadArrayAs<ImportFunctionDescriptor32>().Cast<IImportFunctionDescriptor>().ToArray();
            }
            if (peModule.ModuleType == PeModuleType.Pe32Plus)
            {
                functionDescriptors = reader.ReadArrayAs<ImportFunctionDescriptor64>().Cast<IImportFunctionDescriptor>().ToArray();
            }
            if (functionDescriptors == null) return;

            for (var index = 0; index < functionDescriptors.Length; index++)
            {
                var importFunctionDescriptor = functionDescriptors[index];

                var thunkRva = lookupTablePosition.Value + (uint) index*importFunctionDescriptor.Size;
                var funcOrdinal = -1;
                string funcName = null;
                if (importFunctionDescriptor.ImportByOrdinal)
                {
                    funcOrdinal = (int) importFunctionDescriptor.Ordinal;
                }
                else
                {
                    var importByNameDataPosition = peModule.RvaToFileOffset(importFunctionDescriptor.HintNameTableRva);
                    if (!importByNameDataPosition.HasValue) continue;

                    reader.SetPosition(importByNameDataPosition.Value);
                    var importByNameData = reader.ReadAs<ImportByNameData>();
                    funcOrdinal = importByNameData.Hint;
                    funcName = importByNameData.Name;
                }

                var importedFunction = new ImportedFunction(funcOrdinal, funcName, thunkRva);
                importedModule.Functions.Add(importedFunction);
            }
        }
    }
}